#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Nov  4 15:05:09 2017

@author: shirhe-lyh
"""
import time

import cv2
import numpy as np
import tensorflow as tf
from PIL import Image



# --------------Model preparation----------------
# Path to frozen detection graph. This is the actual model that is used for
# the object detection.
PATH_TO_CKPT = 'frozen_inference_graph.pb'
#PATH_TO_CKPT = 'saved_model/saved_model.pb'

# Load a (frozen) Tensorflow model into memory
detection_graph = tf.Graph()
with detection_graph.as_default():
    od_graph_def = tf.GraphDef()
    with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
        serialized_graph = fid.read()
        od_graph_def.ParseFromString(serialized_graph)
        tf.import_graph_def(od_graph_def, name='')

image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
# Each box represents a part of the image where a particular
# object was detected.
gboxes = detection_graph.get_tensor_by_name('detection_boxes:0')
# Each score represent how level of confidence for each of the objects.
# Score is shown on the result image, together with the class label.
gscores = detection_graph.get_tensor_by_name('detection_scores:0')
gclasses = detection_graph.get_tensor_by_name('detection_classes:0')
gnum_detections = detection_graph.get_tensor_by_name('num_detections:0')

import os

# TODO: Add class names showing in the image
def detect_image_objects(image_path, sess, detection_graph, logfile):
    # Expand dimensions since the model expects images to have
    # shape: [1, None, None, 3]
    print(image_path, file=logfile)
    image = np.array(Image.open(image_path))
    image_np_expanded = np.expand_dims(image, axis=0)

    # Actual detection.
    (boxes, scores, classes, num_detections) = sess.run(
        [gboxes, gscores, gclasses, gnum_detections],
        feed_dict={image_tensor: image_np_expanded})

    # Visualization of the results of a detection.
    boxes = np.squeeze(boxes)
    scores = np.squeeze(scores)
    height, width = image.shape[:2]
    for i in range(boxes.shape[0]):
        if (scores is None or
                scores[i] > 0.9):
            ymin, xmin, ymax, xmax = boxes[i]
            ymin = int(ymin * height)
            ymax = int(ymax * height)
            xmin = int(xmin * width)
            xmax = int(xmax * width)

            score = None if scores is None else scores[i]
            font = cv2.FONT_HERSHEY_SIMPLEX
            text_x = np.max((0, xmin - 10))
            text_y = np.max((0, ymin - 10))
            cv2.putText(image, 'Detection score: ' + str(score),
                        (text_x, text_y), font, 0.4, (0, 255, 0))
            cv2.rectangle(image, (xmin, ymin), (xmax, ymax),
                          (0, 255, 0), 2)
            print('{}:{}:({} {} {} {})'.format(image_path, score, xmin, ymin, xmax, ymax))
            print('{}:{}:({} {} {} {})'.format(image_path, score, xmin, ymin, xmax, ymax), file=logfile)
            p, n = os.path.split(image_path)
            cv2.imwrite('../d03o/' + n, image)
    return image

import pathlib
import sys

PATH_TO_TEST_IMAGES_DIR = pathlib.Path('../d03')
if len(sys.argv) > 1:
  PATH_TO_TEST_IMAGES_DIR = sys.argv[1]
TEST_IMAGE_PATHS = sorted(list(PATH_TO_TEST_IMAGES_DIR.glob("*.jpg")))
OUTPUT = 'output.txt'


with detection_graph.as_default():
    with tf.Session(graph=detection_graph) as sess:
        print('start')
        with open(OUTPUT, 'w') as logfile:
          for image_path in TEST_IMAGE_PATHS:
            print(image_path)
            #image = cv2.imread(image_path)
            detect_image_objects(image_path, sess, detection_graph, logfile)
